import matplotlib.pyplot as plt
import torch
from PIL import Image
from torchvision import transforms
from torchvision import models
from torch import nn
from pytorch_grad_cam import GradCAM, ScoreCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import numpy as np
from args import get_args
import seaborn as sns


def swin_reshape_transform(tensor):
    # Bring the channels to the first dimension,
    # like in CNNs.
    result = tensor.transpose(2, 3).transpose(1, 2)
    return result


def vit_reshape_transform(tensor, height=14, width=14):
    result = tensor[:, 1:, :].reshape(tensor.size(0), height, width, tensor.size(2))

    # Bring the channels to the first dimension,
    # like in CNNs.
    result = result.transpose(2, 3).transpose(1, 2)
    return result


def heat_map(img_path, model_name, model, cam_type):
    img = Image.open(img_path)
    img_size = 224 if model_name == 'vit' else 448

    data_transforms = transforms.Compose([
        transforms.Resize(int(img_size / 0.875)),
        transforms.CenterCrop(img_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # Preprocess it for your chosen model
    input_tensor = data_transforms(img)

    if model_name == 'resnet':
        target_layers = [model.layer4[-1]]
        reshape_transform = None
    elif model_name == 'vit':
        target_layers = [model.encoder.layers[-1].ln_1]
        reshape_transform = vit_reshape_transform
    elif model_name == 'swin':
        target_layers = [model.norm]
        reshape_transform = swin_reshape_transform
    else:
        raise ValueError(f'no model_name called {model_name}')

    if cam_type == 'gradcam':
        cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True,
                      reshape_transform=reshape_transform)
    elif cam_type == 'scorecam':
        cam = ScoreCAM(model=model, target_layers=target_layers, use_cuda=True,
                       reshape_transform=reshape_transform)
    else:
        raise ValueError(f'no such cam named {cam_type}!')

    grayscale_cam = cam(input_tensor=input_tensor.unsqueeze(0))
    grayscale_cam = grayscale_cam[0, :]

    mean = torch.FloatTensor([0.485, 0.456, 0.406]).view(3, 1, 1).expand(3, img_size, img_size)
    std = torch.FloatTensor([0.229, 0.224, 0.225]).view(3, 1, 1).expand(3, img_size, img_size)
    input_tensor = input_tensor * std + mean
    input_tensor = input_tensor.permute(1, 2, 0)
    cam_image = show_cam_on_image(np.asarray(input_tensor), grayscale_cam, use_rgb=True, image_weight=0.5)

    if img_size == 224:
        return Image.fromarray(cam_image).resize((448, 448))
    else:
        return Image.fromarray(cam_image)


def get_cam_pics(model_name, class_nums, weight_paths, imgs_path, cam_type):
    results = []
    if model_name == 'resnet50':
        resnet = models.resnet50(weights=None).eval()
        in_channel = resnet.fc.in_features
        for img_path, resnet_weight, class_num in zip(imgs_path, weight_paths, class_nums):
            resnet.fc = nn.Linear(in_channel, class_num)
            resnet.load_state_dict(torch.load(resnet_weight))
            results.append(heat_map(img_path, 'resnet', resnet, cam_type))

    elif model_name in ['vit-b-16', 'vit-l-16']:
        if model_name == 'vit-b-16':
            vit = models.vit_b_16(weights=None).eval()
        else:
            vit = models.vit_l_16(weights=None).eval()
        in_channel = vit.heads.head.in_features
        for img_path, vit_weight, class_num in zip(imgs_path, weight_paths, class_nums):
            vit.heads.head = nn.Linear(in_channel, class_num)
            vit.load_state_dict(torch.load(vit_weight))
            results.append(heat_map(img_path, 'vit', vit, cam_type))

    elif model_name in ['swin-t', 'swin-b']:
        if model_name == 'swin-t':
            swin = models.swin_t(weights=None).eval()
        else:
            swin = models.swin_b(weights=None).eval()
        in_channel = swin.head.in_features
        for img_path, swin_weight, class_num in zip(imgs_path, weight_paths, class_nums):
            swin.head = nn.Linear(in_channel, class_num)
            swin.load_state_dict(torch.load(swin_weight))
            results.append(heat_map(img_path, 'swin', swin, cam_type))
    return results


def draw_cam(args):
    if len(args.imgs_path) != len(args.datasets):
        raise ValueError('the length of imgs_path must be equal of datasets')

    class_nums = []
    for dataset in args.datasets:
        if dataset == 'cub':
            class_nums.append(200)
        elif dataset == 'dogs':
            class_nums.append(120)
        elif dataset == 'airs':
            class_nums.append(100)

    results = []
    for model in args.models:
        if model == 'resnet50':
            results.append(get_cam_pics(model, class_nums, args.resnet_path, args.imgs_path, 'scorecam'))
        elif model == 'vit-b-16':
            results.append(get_cam_pics(model, class_nums, args.vitb_path, args.imgs_path, 'scorecam'))
        elif model == 'vit-l-16':
            results.append(get_cam_pics(model, class_nums, args.vitl_path, args.imgs_path, 'scorecam'))
        elif model == 'swin-t':
            results.append(get_cam_pics(model, class_nums, args.swint_path, args.imgs_path, 'scorecam'))
        elif model == 'swin-b':
            results.append(get_cam_pics(model, class_nums, args.swinb_path, args.imgs_path, 'gradcam'))

    return results


if __name__ == '__main__':
    args = get_args()

    font_size = 16
    dpi = 100
    params = {'axes.titlesize': font_size,
              'legend.fontsize': font_size,
              'figure.dpi': dpi,
              'figure.figsize': (args.fig_size[0], args.fig_size[1]),
              'axes.labelsize': font_size,
              'xtick.labelsize': font_size,
              'ytick.labelsize': font_size,
              'figure.titlesize': font_size}
    plt.rcParams.update(params)
    plt.rcParams["font.sans-serif"] = ["Times New Roman"]  # 设置字体
    plt.rcParams["axes.unicode_minus"] = False  # 该语句解决图像中的“-”负号的乱码问题
    plt.style.use('seaborn-whitegrid')
    sns.set_style("white")

    results = draw_cam(args)

    model_nums = len(args.models)
    imgs_nums = len(args.imgs_path)
    if model_nums == 1 and imgs_nums == 1:
        plt.imshow(results[0][0])
        plt.axis('off')
    else:
        fig, axes = plt.subplots(nrows=model_nums, ncols=imgs_nums)
        for i in range(model_nums):
            for j in range(imgs_nums):
                axes.flat[imgs_nums * i + j].imshow(results[i][j])
                axes.flat[imgs_nums * i + j].axis('off')
    plt.tight_layout()
    plt.savefig('test.jpg')
    print('Done!')
